package com.syf.starter.core.netty.tool.spi.impl;

import com.syf.starter.core.netty.tool.spi.SyfServerInvokerShape;

import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 平滑漏桶限流器
 * */
public class LeakBucketShape implements SyfServerInvokerShape {
    private final int capacity; // 桶的容量
    private final int leakRate; // 漏水速率（每秒处理的请求数）
    private AtomicInteger water; // 当前桶中的水量
    private final ScheduledExecutorService scheduler;

    private ConcurrentLinkedQueue<CountDownLatch> waitQueue = new ConcurrentLinkedQueue<CountDownLatch>();

    public LeakBucketShape(Integer threadNum, Float avgResponseSeconds) {
        this.capacity = threadNum;
        this.leakRate = setRate(threadNum, avgResponseSeconds).intValue();
        this.water = new AtomicInteger(0);
        this.scheduler = Executors.newScheduledThreadPool(1);
        // 定时任务，每秒执行一次漏水操作
        this.scheduler.scheduleAtFixedRate(this::leak, 0, 1, TimeUnit.SECONDS);
    }

    @Override
    public boolean tryAcquire(long timeout, TimeUnit unit) {
        long timeStamp = System.currentTimeMillis();

        if (water.get() < capacity) {
            water.incrementAndGet();
            return true; // 请求被接受
        } else {
            CountDownLatch latch = new CountDownLatch(1);
            waitQueue.add(latch);

            try {
                if (latch.await(unit.toMillis(timeout) - System.currentTimeMillis() + timeStamp, TimeUnit.MILLISECONDS)) {
                    return true;
                }
            } catch (InterruptedException e) {
                return false;
            }

            return false;
        }
    }

    @Override
    public void acquire() {
        throw new RuntimeException("漏桶流量整形器不支持该方法！");
    }

    @Override
    public void shutdown() {
        scheduler.shutdown();
    }

    // 漏水操作
    private void leak() {
        int currentWater = water.get();
        int newWater = Math.max(0, currentWater - leakRate);
        water.set(newWater);

        int size = waitQueue.size();
        for (int i = 0; i < size; i++) {
            CountDownLatch latch = waitQueue.poll();
            latch.countDown();
        }
    }
}
